"""
智谱清言API客户端封装
"""
import requests
import json
from typing import Dict, List, Optional, BinaryIO, Iterator, Union
import os


class ZhipuAPI:
    def __init__(self, api_key: str, base_url: str = 'https://open.bigmodel.cn/api/paas/v4',
                 upload_folder: str = 'uploads', audio_folder: str = 'audio', 
                 generated_folder: str = 'generated'):
        self.api_key = api_key
        self.base_url = base_url
        self.upload_folder = upload_folder
        self.audio_folder = audio_folder
        self.generated_folder = generated_folder
        
        # 确保目录存在
        os.makedirs(upload_folder, exist_ok=True)
        os.makedirs(audio_folder, exist_ok=True)
        os.makedirs(generated_folder, exist_ok=True)
        
        self.headers = {
            'Authorization': f'Bearer {api_key}',
            'Content-Type': 'application/json'
        }
    
    def _resolve_file_path(self, file_path: str) -> str:
        """解析文件路径，优先在项目目录中查找"""
        # 如果是绝对路径，直接返回
        if os.path.isabs(file_path):
            return file_path
        
        # 如果文件存在，直接返回
        if os.path.exists(file_path):
            return file_path
        
        # 在常用目录中查找
        search_dirs = [
            self.upload_folder,  # 上传的文件
            self.audio_folder,   # 音频文件
            self.generated_folder,  # 生成的文件
            '.',  # 当前目录
        ]
        
        for search_dir in search_dirs:
            full_path = os.path.join(search_dir, file_path)
            if os.path.exists(full_path):
                return full_path
        
        # 如果都找不到，返回原始路径（让调用者处理错误）
        return file_path

    def _request(self, method: str, endpoint: str, **kwargs) -> Dict:
        """通用请求方法"""
        url = f"{self.base_url}{endpoint}"
        try:
            if method.upper() == 'GET':
                response = requests.get(url, headers=self.headers, **kwargs)
            elif method.upper() == 'POST':
                response = requests.post(url, headers=self.headers, **kwargs)
            else:
                return {'success': False, 'error': f'不支持的HTTP方法: {method}'}

            response.raise_for_status()
            return {'success': True, 'data': response.json()}
        except requests.exceptions.RequestException as e:
            return {'success': False, 'error': str(e)}

    def _request_file(self, method: str, endpoint: str, files: dict, data: dict = None) -> Dict:
        """文件上传请求方法"""
        url = f"{self.base_url}{endpoint}"
        headers = {'Authorization': f'Bearer {self.api_key}'}
        try:
            response = requests.post(url, headers=headers, files=files, data=data)
            response.raise_for_status()
            # 尝试解析JSON响应
            try:
                response_data = response.json()
                return {'success': True, 'data': response_data}
            except ValueError:
                # 如果响应不是JSON，返回文本内容
                return {'success': True, 'data': {'text': response.text}}
        except requests.exceptions.HTTPError as e:
            # HTTP错误，尝试获取错误详情
            try:
                error_data = e.response.json() if e.response else {}
                error_msg = error_data.get('error', {}).get('message', str(e)) if isinstance(error_data.get('error'), dict) else error_data.get('message', str(e))
                return {'success': False, 'error': error_msg, 'status_code': e.response.status_code if e.response else None}
            except:
                return {'success': False, 'error': f'HTTP错误: {e.response.status_code if e.response else "未知"} - {str(e)}'}
        except requests.exceptions.RequestException as e:
            return {'success': False, 'error': str(e)}

    # ============ 对话补全 ============
    def chat_completion(self, messages: List[Dict], model: str = 'glm-4.5-flash',
                       temperature: float = 1, max_tokens: int = 65536) -> Iterator[Dict]:
        """对话补全（流式输出）"""
        endpoint = '/chat/completions'
        url = f"{self.base_url}{endpoint}"
        payload = {
            'model': model,
            'messages': messages,
            'temperature': temperature,
            'max_tokens': max_tokens,
            'stream': True,
            'thinking': {
                'type': 'disabled'
            }
        }
        
        try:
            response = requests.post(url, headers=self.headers, json=payload, stream=True)
            response.raise_for_status()
            
            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith('data: '):
                        data_str = line_str[6:]  # 移除 'data: ' 前缀
                        if data_str.strip() == '[DONE]':
                            break
                        try:
                            data = json.loads(data_str)
                            yield {'success': True, 'data': data}
                        except json.JSONDecodeError:
                            continue
        except requests.exceptions.RequestException as e:
            yield {'success': False, 'error': str(e)}

    # ============ 图像生成 ============
    def generate_image(self, prompt: str, model: str = 'cogview-3-flash', 
                      size: str = '1024x1024', save_to_file: bool = True, 
                      filename: str = None) -> Dict:
        """生成图像
        
        Args:
            prompt: 图像生成提示词
            model: 模型名称
            size: 图像尺寸
            save_to_file: 是否自动下载并保存到generated目录
            filename: 保存的文件名（如果不指定则自动生成）
        """
        endpoint = '/images/generations'
        payload = {
            'model': model,
            'prompt': prompt,
            'size': size
        }
        response = self._request('POST', endpoint, json=payload)
        
        # 如果成功且启用自动保存，下载图像
        if response.get('success') and save_to_file:
            try:
                data = response['data']
                if 'data' in data and len(data['data']) > 0:
                    image_url = data['data'][0].get('url')
                    if image_url:
                        # 下载图像
                        img_response = requests.get(image_url)
                        img_response.raise_for_status()
                        
                        if not filename:
                            import time
                            filename = f"image_{int(time.time())}.png"
                        
                        file_path = os.path.join(self.generated_folder, filename)
                        with open(file_path, 'wb') as f:
                            f.write(img_response.content)
                        
                        response['file_path'] = file_path
                        response['filename'] = filename
            except Exception as e:
                # 下载失败不影响主要结果
                response['download_error'] = str(e)
        
        return response

    # ============ 视频生成 ============
    def generate_video(self, prompt: str, model: str = 'cogvideox-flash',
                      quality: str = 'quality', with_audio: bool = True,
                      size: str = '1920x1080', fps: int = 30) -> Dict:
        """生成视频（异步）"""
        endpoint = '/videos/generations'
        payload = {
            'model': model,
            'prompt': prompt,
            'quality': quality,
            'with_audio': with_audio,
            'size': size,
            'fps': fps
        }
        return self._request('POST', endpoint, json=payload)

    # ============ 查询异步结果 ============
    def get_async_result(self, task_id: str) -> Dict:
        """查询异步结果"""
        endpoint = f'/async-result/{task_id}'
        return self._request('GET', endpoint)
    
    def download_video(self, video_url: str, filename: str = None, save_to_file: bool = True) -> Dict:
        """下载视频文件
        
        Args:
            video_url: 视频URL
            filename: 保存的文件名（如果不指定则自动生成）
            save_to_file: 是否保存到generated目录
        """
        try:
            # 下载视频
            video_response = requests.get(video_url, stream=True)
            video_response.raise_for_status()
            
            result = {'success': True}
            
            if save_to_file:
                if not filename:
                    import time
                    # 从URL获取文件扩展名，默认使用mp4
                    ext = 'mp4'
                    if '.' in video_url:
                        url_ext = video_url.rsplit('.', 1)[1].lower().split('?')[0]
                        if url_ext in ['mp4', 'webm', 'mov']:
                            ext = url_ext
                    filename = f"video_{int(time.time())}.{ext}"
                
                file_path = os.path.join(self.generated_folder, filename)
                # 使用流式下载保存大文件
                with open(file_path, 'wb') as f:
                    for chunk in video_response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                
                result['file_path'] = file_path
                result['filename'] = filename
            
            return result
        except requests.exceptions.RequestException as e:
            return {'success': False, 'error': str(e)}

    # 已移除：语音转文本

    # ============ 文本转语音 ============
    def text_to_speech(self, text: str, model: str = 'cogtts', 
                      voice: str = 'tongtong', response_format: str = 'wav',
                      save_to_file: bool = True, filename: str = None) -> Dict:
        """文本转语音
        
        Args:
            text: 要转换的文本
            model: 模型名称
            voice: 语音类型
            response_format: 音频格式（wav, mp3等）
            save_to_file: 是否自动保存到audio目录
            filename: 保存的文件名（如果不指定则自动生成）
        """
        endpoint = '/audio/speech'
        payload = {
            'model': model,
            'input': text,
            'voice': voice,
            'response_format': response_format
        }
        try:
            url = f"{self.base_url}{endpoint}"
            headers = {'Authorization': f'Bearer {self.api_key}'}
            response = requests.post(url, headers=headers, json=payload)
            response.raise_for_status()
            # 文本转语音返回的是二进制音频数据
            audio_data = response.content
            
            result = {'success': True, 'data': audio_data}
            
            # 如果启用自动保存，保存到audio目录
            if save_to_file:
                if not filename:
                    import time
                    filename = f"speech_{int(time.time())}.{response_format}"
                
                file_path = os.path.join(self.audio_folder, filename)
                with open(file_path, 'wb') as f:
                    f.write(audio_data)
                result['file_path'] = file_path
                result['filename'] = filename
            
            return result
        except requests.exceptions.RequestException as e:
            return {'success': False, 'error': str(e)}
 
    # 已移除：网络搜索、文件解析相关方法

